#include "../include/func.h"
#include <openssl/sha.h>

// Epoll 相关函数
int epollAdd(int epfd, int fd) {
    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = fd;
    epoll_ctl(epfd, EPOLL_CTL_ADD, fd, &ev);
    return 0;
}

int epollDel(int epfd, int fd) {
    epoll_ctl(epfd, EPOLL_CTL_DEL, fd, NULL);
    return 0;
}

// 配置读取函数
int getValidStr(char *input, char *str) {
    char *p; 
    char s[] = {' ', '\n'};
    p = strtok(input, s);
    p = strtok(NULL, s);
    p = strtok(NULL, s);
    if (p != NULL) {
        strcpy(str, p); 
    }
    return 0;
}

int getClientConf(char *fileconf, char *ip, char *port) {
    FILE *pconf = fopen(fileconf, "r");
    ERROR_CHECK(pconf, NULL, "open conf file failed");
    char buf[128] = {0};
    fgets(buf, sizeof(buf), pconf);
    getValidStr(buf, ip);
    bzero(buf, sizeof(buf));
    fgets(buf, sizeof(buf), pconf);
    getValidStr(buf, port);
    fclose(pconf);
    return 0;
}

// 命令解析函数
void init_cmd(char *line, pcmd_t pcmd) {
    char *token;
    const char s[] = " \n";

    token = strtok(line, s);
    while (NULL != token && pcmd->argc < MAX_WORDNUM) {
        strncpy(pcmd->cmd[pcmd->argc], token, MAX_WORDLEN - 1);
        pcmd->cmd[pcmd->argc][MAX_WORDLEN - 1] = '\0';
        pcmd->argc++;
        token = strtok(NULL, s);
    }
}

void print_cmd(pcmd_t pcmd) {
    printf("命令解析: ");
    for (int i = 0; i < pcmd->argc; ++i) {
        printf("%s ", pcmd->cmd[i]);
    }
    printf("\n");
}

int cmd_check(pcmd_t pcmd, int *ptag) {
    int flag = -1;
    char tmp[64] = "";
    strncpy(tmp, pcmd->cmd[0], 64);

    if (strcmp(tmp, "cd") == 0) {
        *ptag = TAGCD;
        flag = 0;
    } else if (strcmp(tmp, "ls") == 0) {
        *ptag = TAGLS;
        flag = 0;
    } else if (strcmp(tmp, "pwd") == 0) {
        *ptag = TAGPWD;
        flag = 0;
    } else if (strcmp(tmp, "rm") == 0) {
        *ptag = TAGREMOVE;
        flag = 0;
    } else if (strcmp(tmp, "puts") == 0) {
        *ptag = TAGPUTS;
        flag = 0;
    } else if (strcmp(tmp, "gets") == 0) {
        *ptag = TAGGETS;
        flag = 0;
    } else if (strcmp(tmp, "mkdir") == 0) {
        *ptag = TAGMKDIR;
        flag = 0;
    } else if (strcmp(tmp, "rmdir") == 0) {
        *ptag = TAGRMDIR;
        flag = 0;
    }
    return flag;
}

int user_register(socket_t fd_server, userInfo_t *user_info) {
    train_t train;
    char username[256] = {0};
    char password[256] = {0};
    char salt[256] = {0};

    //  发送注册标志和用户名
    train.tag = TAG_REGISTER;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 读取用户的用户名输入
    // 不允许用户输入空字符
    do {
        printf("请输入用户名：");
        scanf("%s", username);
        getchar(); // 清除换行符
    } while (!strcmp(username,""));

    // 发送用户名
    train.data_length = strlen(username);
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, username, train.data_length, 0);

    //接收盐值
    recv(fd_server, &train.data_length, sizeof(int), 0); 
    if (train.data_length == 0) {
        printf("该用户名已经被注册!\n");
        return -1;
    }

    recv(fd_server, salt, train.data_length, MSG_WAITALL);
    printf("收到盐值: %s\n", salt);

    do {
        printf("请输入密码: ");
        scanf("%s", password);   // 输入回车/EOF/纯空格的情况
        getchar(); // 清除换行符
    } while (!strcmp(password,""));

    //  用盐值和明文密码加密得到新密文
    char *encrypted_password = crypt(password, salt);
    if (!encrypted_password) {
        printf("错误: 密码加密失败\n");
        return -1;
    }

    printf("加密后的密码: %s\n", encrypted_password);

    // 发送加密后的密码
    train.data_length = strlen(encrypted_password);
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, encrypted_password, train.data_length, 0);

    // 接收最终注册结果
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);
    if (train.tag == TAG_ERROR) {
        printf("%s\n", train.data);
        return -1;
    }

    printf("%s\n", train.data);

    // 更新 userInfo 
    strcpy(user_info->username, username);
    strcpy(user_info->current_dir,"/");
    strcat(user_info->current_dir,username);

    return 0;
}

int user_login(socket_t fd_server, userInfo_t *userInfo) {
    train_t train;
    char username[256] = {0};
    char password[256] = {0};
    char salt[256] = {0};
    // 发送登录请求和用户名
    train.tag = TAG_LOGIN;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 读取用户的用户名输入
    // 不允许用户输入空字符
    do {
        printf("请输入用户名：");
        scanf("%s", username);
        getchar(); // 清除换行符
    } while (!strcmp(username,""));

    train.data_length = strlen(username);
    send(fd_server, &train.data_length, sizeof(int), 0); 
    send(fd_server, username, train.data_length, 0);

    printf("%s\n",username);
    recv(fd_server,&train.data_length,sizeof(int),0);
    if(train.data_length == 0){
        printf("该用户不存在！\n");
        return -1;
    }
    recv(fd_server, salt, train.data_length, MSG_WAITALL);

    printf("请输入密码: ");
    scanf("%s", password);
    getchar(); // 清除换行符

    // 加密密码
    char *encrypted_password = crypt(password, salt);

    // 发送加密后的密码
    train.data_length = strlen(encrypted_password);
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, encrypted_password, train.data_length, 0);

    // 接收登录结果
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);

    if (train.tag == TAG_RESPONSELOGIN) {
        printf("%s\n", train.data);
        // 登录成功，更新共享资源
        strcpy(userInfo->username, username);
        strcpy(userInfo->current_dir,"/");
        strcat(userInfo->current_dir,username);
        /* userInfo->user_id = atoi(train.data); */
        /* strcpy(userInfo->username, username); */
        /* printf("登录成功! 用户ID: %d\n", userInfo->user_id); */
        // 接收当前目录
        /* recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL); */
        /* recv(fd_server, userInfo->current_dir, train.data_length, MSG_WAITALL); */
        /* printf("当前目录: %s\n", userInfo->current_dir); */
        return 0;
    } else {
        printf("登录失败: 密码错误\n");
        return -1;
    }
}
int user_cancel(socket_t fd_server, userInfo_t *userInfo) {
    train_t train;
    char username[256] = {0};
    char password[256] = {0};
    char salt[256] = {0};

    // 发送注销标志
    train.tag = TAG_CANCEL;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 读取用户的用户名输入
    // 不允许用户输入空字符
    do {
        printf("请输入用户名：");
        scanf("%s", username);
        getchar(); // 清除换行符
    } while (!strcmp(username,""));

    // 发送用户名
    train.data_length = strlen(username);
    send(fd_server, &train.data_length, sizeof(int), 0); 
    send(fd_server, username, train.data_length, 0); 

    //接收盐值
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);                    
    if (train.data_length == 0) {
        printf("该用户不存在!\n");
        return -1; 
    }   

    recv(fd_server, salt, train.data_length, MSG_WAITALL);
    printf("收到盐值: %s\n", salt);

    printf("请输入密码: ");
    scanf("%s", password);
    getchar(); // 清除换行符

    // 用盐值和明文密码加密得到新密文
    char *encrypted_password = crypt(password, salt);
    if (!encrypted_password) {
        printf("错误: 密码加密失败\n");
        return -1; 
    }   

    printf("加密后的密码: %s\n", encrypted_password);

    // 发送加密后的密码
    train.data_length = strlen(encrypted_password);
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, encrypted_password, train.data_length, 0);

    // 接收响应
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    if(train.tag != TAG_RESPONSECANCEL){
        printf("用户[%s]身份验证失败>_<！\n", username);
        return -1;
    }else{
        printf("用户[%s]已注销>_<\n", username);
        return 0;
    }

    (void)userInfo;
}

// 目录操作函数
int handle_pwd(socket_t fd_server) {
    train_t train;
    train.tag = TAGPWD;
    send(fd_server, &train.tag, sizeof(int), 0);

    bzero(&train, sizeof(train_t));
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);

    printf("当前路径: %s\n", train.data);
    return 0;
}

int handle_ls(socket_t fd_server) {
    train_t train;
    train.tag = TAGLS;
    send(fd_server, &train.tag, sizeof(int), 0);

    bzero(&train, sizeof(train_t));
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);
    if(train.tag == TAG_ERROR){
        printf("%s",train.data);
        return -1;
    }
    printf("当前目录:%s", train.data);
    // 对面会循环发送目录项
    return 0;
}

int handle_cd(socket_t fd_server, pcmd_t pcmd) {
    if (pcmd->argc < 2) {
        printf("用法: cd <目录>\n");
        return -1;
    }

    train_t train;
    train.tag = TAGCD;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 发送目录路径
    train.data_length = strlen(pcmd->cmd[1]) + 1;
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, pcmd->cmd[1], train.data_length, 0);

    // 接收结果
    bzero(&train, sizeof(train_t));
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);
    if(train.tag == TAG_ERROR){
        printf("路径跳转错误：%s\n", train.data);
        return -1;
    }
    printf("当前路径: %s\n", train.data);
    return 0;    
}

int handle_mkdir(socket_t fd_server, pcmd_t pcmd) {
    if (pcmd->argc < 2) {
        printf("用法: mkdir <目录名>\n");
        return -1;
    }
    train_t train;
    train.tag = TAGMKDIR;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 发送目录名
    train.data_length = strlen(pcmd->cmd[1]) + 1;
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, pcmd->cmd[1], train.data_length, 0);

    // 接收结果
    bzero(&train, sizeof(train_t));
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);

    printf("结果: %s\n", train.data);
    return 0;
}

int handle_rmdir(socket_t fd_server, pcmd_t pcmd) {
    if (pcmd->argc < 2) {
        printf("用法: rmdir <目录名>\n");
        return -1;
    }
    train_t train;
    train.tag = TAGRMDIR;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 发送目录名
    train.data_length = strlen(pcmd->cmd[1]) + 1;
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, pcmd->cmd[1], train.data_length, 0);

    // 接收结果
    bzero(&train, sizeof(train_t));
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);

    printf("结果: %s\n", train.data);
    return 0;
}

int handle_remove(socket_t fd_server, pcmd_t pcmd) {
    if (pcmd->argc < 2) {
        printf("用法: rm <文件名>\n");
        return -1;
    }
    train_t train;
    train.tag = TAGREMOVE;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 发送文件名
    train.data_length = strlen(pcmd->cmd[1]) + 1;
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, pcmd->cmd[1], train.data_length, 0);

    // 接收结果
    bzero(&train, sizeof(train_t));
    recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);

    printf("结果: %s\n", train.data);
    return 0;
}

// 修改客户端的哈希计算函数，使用与服务端相同的 EVP 函数
char* calculate_file_hash(char *filename) {
    FILE *file = fopen(filename, "rb");
    if (!file) {
        printf("错误: 无法打开文件 '%s' 计算哈希值\n", filename);
        return NULL;
    }

    EVP_MD_CTX *mdctx = EVP_MD_CTX_new();
    if (!mdctx) {
        fclose(file);
        return NULL;
    }

    // 使用现代API初始化SHA512，与服务端保持一致
    if (EVP_DigestInit_ex(mdctx, EVP_sha512(), NULL) != 1) {
        EVP_MD_CTX_free(mdctx);
        fclose(file);
        return NULL;
    }

    unsigned char buffer[1024];
    size_t bytesRead;
    while ((bytesRead = fread(buffer, 1, sizeof(buffer), file)) != 0) {
        if (EVP_DigestUpdate(mdctx, buffer, bytesRead) != 1) {
            EVP_MD_CTX_free(mdctx);
            fclose(file);
            return NULL;
        }
    }

    unsigned char hash[EVP_MAX_MD_SIZE];
    unsigned int hash_length;
    if (EVP_DigestFinal_ex(mdctx, hash, &hash_length) != 1) {
        EVP_MD_CTX_free(mdctx);
        fclose(file);
        return NULL;
    }

    EVP_MD_CTX_free(mdctx);
    fclose(file);

    // 将哈希值转换为十六进制字符串，与服务端格式一致
    char *hash_str = (char *)malloc(hash_length * 2 + 1);
    if (!hash_str) {
        return NULL;
    }

    for (unsigned int i = 0; i < hash_length; i++) {
        sprintf(hash_str + (i * 2), "%02x", hash[i]);
    }
    hash_str[hash_length * 2] = '\0';

    return hash_str;
}

// 修改发送文件函数，与服务端协议匹配
int send_file_train(int socket_fd, char *filename, off_t file_size) {
    int file_fd = open(filename, O_RDONLY);
    ERROR_CHECK(file_fd, -1, "open");

    train_t train;
    bzero(&train, sizeof(train_t));
    int ret = -1;
    off_t total_sent = 0;

    printf("文件大小: %ld bytes (%.2f MB)\n", file_size, (double)file_size / (1024 * 1024));

    // 先发送文件大小 - 与服务端匹配
    train.data_length = sizeof(off_t);
    send(socket_fd, &train.data_length, sizeof(int), 0);
    send(socket_fd, &file_size, train.data_length, 0);

    // 接收服务端的断点偏移量
    off_t resume_offset = 0;
    recv(socket_fd, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(socket_fd, &resume_offset, train.data_length, MSG_WAITALL);

    printf("服务端续传偏移量: %ld bytes\n", resume_offset);

    // 如果服务端已经有部分文件，定位到续传位置
    if (resume_offset > 0) {
        off_t seek_ret = lseek(file_fd, resume_offset, SEEK_SET);
        if (seek_ret == (off_t)-1) {
            printf("定位文件断点失败\n");
            close(file_fd);
            return -1;
        }
        total_sent = resume_offset;
        printf("从断点开始续传，已跳过 %ld bytes\n", resume_offset);
    }

    // 大文件检测：超过100MB使用mmap零拷贝
    if (file_size > 100 * 1024 * 1024) {
        printf("大文件检测，使用mmap零拷贝传输...\n");
        printf("提示: 传输过程中按 'q' 可中断传输\n");

        // 使用mmap映射文件
        void *file_map = mmap(NULL, file_size, PROT_READ, MAP_PRIVATE, file_fd, 0);
        if (file_map == MAP_FAILED) {
            printf("mmap失败，使用普通传输: %s\n", strerror(errno));
            lseek(file_fd, resume_offset, SEEK_SET);
        } else {
            // mmap传输逻辑
            const off_t CHUNK_SIZE = sizeof(train_t);
            off_t current_offset = resume_offset;

            while (current_offset < file_size) {
                off_t remaining = file_size - current_offset;
                off_t chunk_size = (remaining > CHUNK_SIZE) ? CHUNK_SIZE : remaining;

                // 发送数据块大小
                train.data_length = chunk_size;
                ret = send(socket_fd, &train.data_length, sizeof(int), 0);
                if (ret == -1) {
                    printf("连接异常断开1！\n");
                    break;
                }

                // 使用mmap直接发送数据
                ret = send(socket_fd, (char*)file_map + current_offset, chunk_size, 0);
                if (ret == -1) {
                    printf("连接异常断开2！\n");
                    break;
                }

                current_offset += chunk_size;
                total_sent = current_offset;

                // 显示进度
                int percent = (int)((double)total_sent / file_size * 100);
                printf("\r传输进度: %d%% (%ld/%ld bytes)", percent, total_sent, file_size);
                fflush(stdout);

                // 检查用户中断
                fd_set readfds;
                struct timeval timeout;
                FD_ZERO(&readfds);
                FD_SET(STDIN_FILENO, &readfds);
                timeout.tv_sec = 0;
                timeout.tv_usec = 1000;

                if (select(STDIN_FILENO + 1, &readfds, NULL, NULL, &timeout) > 0) {
                    char cmd;
                    read(STDIN_FILENO, &cmd, 1);
                    if (cmd == 'q' || cmd == 'Q') {
                        printf("\n用户中断传输，已传输: %ld bytes\n", total_sent);
                        break;
                    }
                }
            }

            munmap(file_map, file_size);

            if (current_offset == file_size) {
                // 发送传输完成标志
                train.data_length = 0;
                send(socket_fd, &train.data_length, sizeof(int), 0);
                printf("\n文件传输完成\n");
                close(file_fd);
                return 0;
            } else {
                printf("\n文件传输中断\n");
                close(file_fd);
                return -2;
            }
        }
    }

    // 普通文件传输
    while (1) {
        memset(train.data, 0, sizeof(train.data));
        train.data_length = read(file_fd, train.data, sizeof(train.data));
        ERROR_CHECK(train.data_length, -1, "read");

        if (train.data_length == 0) {
            // 传输完成
            ret = send(socket_fd, &train.data_length, sizeof(int), 0);
            ERROR_CHECK(ret, -1, "send");
            close(file_fd);
            printf("\n文件传输完成\n");
            return 0;
        }

        // 发送数据长度
        ret = send(socket_fd, &train.data_length, sizeof(int), 0);
        if (ret == -1) {
            printf("连接异常断开！\n");
            close(file_fd);
            return -2;
        }

        // 发送数据内容
        ret = send(socket_fd, train.data, train.data_length, 0);
        if (ret == -1) {
            printf("连接异常断开！\n");
            close(file_fd);
            return -2;
        }

        total_sent += train.data_length;
        // 显示传输进度
        if (file_size > 0) {
            int percent = (int)((double)total_sent / file_size * 100);
            printf("\r传输进度: %d%% (%ld/%ld bytes)", percent, total_sent, file_size);
            fflush(stdout);
        }

        // 检查用户中断
        fd_set readfds;
        struct timeval timeout;
        FD_ZERO(&readfds);
        FD_SET(STDIN_FILENO, &readfds);
        timeout.tv_sec = 0;
        timeout.tv_usec = 1000;

        if (select(STDIN_FILENO + 1, &readfds, NULL, NULL, &timeout) > 0) {
            char cmd;
            if (read(STDIN_FILENO, &cmd, 1) > 0) {
                if (cmd == 'q' || cmd == 'Q') {
                    printf("\n用户中断传输，已传输: %ld bytes\n", total_sent);
                    close(file_fd);
                    return -3;
                }
            }
        }
    }

    close(file_fd);
    return 0;
}

// 修改接收文件函数，与服务端协议匹配
int recv_file_train(int fd_socket, char *filename) {
    train_t train;
    int ret = -1;

    // 先接收文件总大小
    off_t file_size = 0;
    recv(fd_socket, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_socket, &file_size, train.data_length, MSG_WAITALL);

    printf("文件总大小: %ld bytes (%.2f MB)\n", file_size, (double)file_size / (1024 * 1024));

    // 检查本地文件是否存在及大小（用于断点续传）
    off_t resume_offset = 0;
    struct stat st;
    if (stat(filename, &st) == 0) {
        resume_offset = st.st_size;
        printf("本地文件已存在，大小: %ld bytes\n", resume_offset);

        if (resume_offset >= file_size) {
            printf("文件已完整下载\n");
            return 0;
        }
    }

    // 发送断点偏移量给服务端 - 与服务端匹配
    train.data_length = sizeof(off_t);
    send(fd_socket, &train.data_length, sizeof(int), 0);
    send(fd_socket, &resume_offset, train.data_length, 0);

    // 大文件检测：超过100MB使用mmap零拷贝
    if (file_size > 100 * 1024 * 1024) {
        printf("大文件检测，使用mmap零拷贝接收...\n");

        int flags = O_WRONLY | O_CREAT;
        if (resume_offset > 0) {
            flags |= O_APPEND;
            printf("从断点 %ld bytes 开始续传\n", resume_offset);
        } else {
            flags |= O_TRUNC;
        }

        int fd_file = open(filename, flags, 0666);
        ERROR_CHECK(fd_file, -1, "open");

        // 预分配文件空间
        if (resume_offset == 0) {
            ftruncate(fd_file, file_size);
        }

        void *file_map = mmap(NULL, file_size, PROT_WRITE, MAP_SHARED, fd_file, 0);
        if (file_map == MAP_FAILED) {
            printf("mmap失败，使用普通接收: %s\n", strerror(errno));
            close(fd_file);
        } else {
            off_t total_received = resume_offset;

            while (total_received < file_size) {
                // 接收数据长度
                ret = recv(fd_socket, &train.data_length, sizeof(int), MSG_WAITALL);
                if (ret <= 0) {
                    printf("连接异常断开\n");
                    break;
                }

                if (train.data_length == 0) {
                    // 传输完成
                    break;
                }

                // 直接接收到mmap映射的内存
                ret = recv(fd_socket, (char*)file_map + total_received, train.data_length, MSG_WAITALL);
                if (ret <= 0) {
                    printf("连接异常断开\n");
                    break;
                }

                total_received += train.data_length;

                // 定期同步到磁盘
                if (total_received % (10 * 1024 * 1024) == 0) {
                    msync(file_map, file_size, MS_ASYNC);
                }

                // 显示进度
                int percent = (int)((double)total_received / file_size * 100);
                printf("\r接收进度: %d%% (%ld/%ld bytes)", percent, total_received, file_size);
                fflush(stdout);
            }

            msync(file_map, file_size, MS_SYNC);
            munmap(file_map, file_size);
            close(fd_file);

            if (total_received == file_size) {
                printf("\n文件接收完成\n");
                return 0;
            } else {
                printf("\n文件接收中断，已接收: %ld bytes\n", total_received);
                return -2;
            }
        }
    }

    // 普通文件接收
    int flags = O_WRONLY | O_CREAT;
    if (resume_offset > 0) {
        flags |= O_APPEND;
        printf("从断点 %ld bytes 开始续传\n", resume_offset);
    } else {
        flags |= O_TRUNC;
    }

    int fd_file = open(filename, flags, 0666);
    ERROR_CHECK(fd_file, -1, "open");

    off_t total_received = resume_offset;

    while (total_received < file_size) {
        memset(&train, 0, sizeof(train));

        // 接收数据长度
        ret = recv(fd_socket, &train.data_length, sizeof(int), MSG_WAITALL);
        if (ret <= 0) {
            printf("连接异常断开\n");
            close(fd_file);
            return -2;
        }

        if (train.data_length == 0) {
            // 传输完成
            break;
        }

        // 接收数据内容
        ret = recv(fd_socket, train.data, train.data_length, MSG_WAITALL);
        if (ret <= 0) {
            printf("连接异常断开\n");
            close(fd_file);
            return -2;
        }

        // 写入文件
        write(fd_file, train.data, train.data_length);
        total_received += train.data_length;

        // 显示接收进度
        if (file_size > 0) {
            int percent = (int)((double)total_received / file_size * 100);
            printf("\r接收进度: %d%% (%ld/%ld bytes)", percent, total_received, file_size);
            fflush(stdout);
        }
    }

    close(fd_file);

    if (total_received == file_size) {
        printf("\n文件接收完成\n");
        return 0;
    } else {
        printf("\n文件接收不完整: %ld/%ld bytes\n", total_received, file_size);
        return -1;
    }
}

// 修改文件上传函数
int handle_puts(socket_t fd_server, pcmd_t pcmd) {
    if (pcmd->argc < 2) {
        printf("用法: puts <客户端文件名> [服务端文件名]\n");
        return -1;
    }

    char *client_filename = pcmd->cmd[1];
    char *server_filename = pcmd->cmd[1]; // 默认使用相同文件名
    
    if (pcmd->argc >= 3) {
        server_filename = pcmd->cmd[2];
    }

    // 检查客户端文件是否存在并获取文件大小
    struct stat st;
    if (stat(client_filename, &st) == -1) {
        printf("错误: 文件 '%s' 不存在\n", client_filename);
        return -1;
    }

    if (!S_ISREG(st.st_mode)) {
        printf("错误: '%s' 不是普通文件\n", client_filename);
        return -1;
    }

    off_t file_size = st.st_size;
    printf("客户端文件 '%s' 存在，大小: %ld bytes (%.2f MB)，准备上传...\n", 
           client_filename, file_size, (double)file_size / (1024 * 1024));

    // 发送 puts tag
    train_t train;
    bzero(&train, sizeof(train_t));
    train.tag = TAGPUTS;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 发送服务端文件名
    train.data_length = strlen(server_filename);
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, server_filename, train.data_length, 0);

    // 计算并发送文件哈希值
    char *file_hash = calculate_file_hash(client_filename);
    if (!file_hash) {
        printf("错误: 无法计算文件哈希值\n");
        train.data_length = 1;
        send(fd_server, &train.data_length, sizeof(int), 0);
        send(fd_server, "", train.data_length, 0);
        return -1;
    }

    printf("文件哈希值: %s\n", file_hash);
    train.data_length = strlen(file_hash);
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, file_hash, train.data_length, 0);
    free(file_hash);

    // 发送文件内容
    printf("开始传输文件 '%s'...\n", client_filename);
    int ret = send_file_train(fd_server, client_filename, file_size);

    if (ret == 0) {
        printf("文件上传成功\n");

        // 接收服务端确认
        recv(fd_server, &train.tag, sizeof(int), MSG_WAITALL);
        recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
        recv(fd_server, train.data, train.data_length, MSG_WAITALL);
        printf("服务端响应: %s\n", train.data);
    } else if (ret == -2) {
        printf("连接异常断开！\n");
    } else if (ret == -3) {
        printf("用户中断传输！\n");
    } else {
        printf("文件上传失败\n");
    }

    return ret;
}

// 修改文件下载函数
int handle_gets(socket_t fd_server, pcmd_t pcmd) {
    if (pcmd->argc < 2) {
        printf("用法: gets <服务端文件名> [客户端文件名]\n");
        return -1;
    }

    char *server_filename = pcmd->cmd[1];
    char *client_filename = pcmd->cmd[1]; // 默认使用相同文件名
    
    if (pcmd->argc >= 3) {
        client_filename = pcmd->cmd[2];
    }

    // 发送 gets tag
    train_t train;
    bzero(&train, sizeof(train_t));
    train.tag = TAGGETS;
    send(fd_server, &train.tag, sizeof(int), 0);

    // 发送服务端文件名
    train.data_length = strlen(server_filename);
    send(fd_server, &train.data_length, sizeof(int), 0);
    send(fd_server, server_filename, train.data_length, 0);

    // 接收服务端响应，检查文件是否存在
    int response_tag;
    recv(fd_server, &response_tag, sizeof(int), MSG_WAITALL);
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, train.data, train.data_length, MSG_WAITALL);

    if (response_tag == TAG_ERROR) {
        printf("%s\n", train.data);
        return -1;
    }

    // 接收服务端传的文件哈希值
    char server_hash[SHA512_DIGEST_LENGTH * 2 + 1] = {0};
    recv(fd_server, &train.data_length, sizeof(int), MSG_WAITALL);
    recv(fd_server, server_hash, train.data_length, MSG_WAITALL);
    printf("服务端文件哈希值: %s\n", server_hash);

    // 接收文件内容
    printf("开始下载文件到 '%s'...\n", client_filename);
    int ret = recv_file_train(fd_server, client_filename);

    if (ret == 0) {
        // 验证文件完整性
        char *client_hash = calculate_file_hash(client_filename);
        if (client_hash && strcmp(client_hash, server_hash) == 0) {
            printf("文件下载成功且完整性验证通过\n");
        } else {
            printf("警告: 文件下载完成但完整性验证失败\n");
            printf("客户端哈希: %s\n", client_hash ? client_hash : "计算失败");
            printf("服务端哈希: %s\n", server_hash);
            ret = -1;
        }
        free(client_hash);
    } else if (ret == -2) {
        printf("连接异常断开，下次可续传\n");
    } else {
        printf("文件下载失败\n");
    }

    return ret;
}
